import requests
import json
from dotenv import load_dotenv
import os
import re
import logging

load_dotenv()

class LLMService:
    def __init__(self):
        self.api_key = os.getenv('DEEPSEEK_API_KEY')
        self.api_url = os.getenv('DEEPSEEK_API_URL')
        self.headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }

    def clean_sql(self, sql):
        """清理SQL语句中的Markdown符号和其他不需要的字符"""
        # 移除Markdown代码块标记
        sql = re.sub(r'```sql|```', '', sql)
        # 移除行首和行尾的空白字符
        sql = sql.strip()
        # 移除SQL语句末尾的分号（如果有）
        sql = re.sub(r';$', '', sql)
        # 移除多余的换行符
        sql = re.sub(r'\n+', ' ', sql)
        # 移除多余的空格
        sql = re.sub(r'\s+', ' ', sql)
        return sql

    def split_compound_query(self, sql):
        """将复合查询拆分为多个独立的SQL语句"""
        # 按分号分割SQL语句
        queries = [q.strip() for q in sql.split(';') if q.strip()]
        return queries

    def generate_sql(self, natural_language_query):
        prompt = f"""请将以下自然语言查询转换为SQL语句。如果查询包含多个操作（如先更新后查询），请用分号分隔每个SQL语句。
        数据库信息：
        - 数据库类型：MySQL
        - 当前数据库：{os.getenv('MYSQL_DATABASE')}
        
        用户查询：{natural_language_query}"""

        data = {
            "model": "deepseek-chat",
            "messages": [
                {"role": "system", "content": "你是一个专业的SQL转换助手，负责将自然语言转换为SQL语句。如果查询包含多个操作，请用分号分隔每个SQL语句。只返回SQL语句，不要包含任何解释或额外文本。"},
                {"role": "user", "content": prompt}
            ],
            "temperature": 0.1
        }

        try:
            response = requests.post(self.api_url, headers=self.headers, json=data)
            response.raise_for_status()
            result = response.json()
            raw_sql = result['choices'][0]['message']['content'].strip()
            
            # 清理SQL语句
            cleaned_sql = self.clean_sql(raw_sql)
            logging.info(f"原始SQL: {raw_sql}")
            logging.info(f"清理后的SQL: {cleaned_sql}")
            
            # 拆分复合查询
            queries = self.split_compound_query(cleaned_sql)
            logging.info(f"拆分后的SQL语句: {queries}")
            
            return queries
        except Exception as e:
            logging.error(f"生成SQL时发生错误: {str(e)}")
            return [f"Error: {str(e)}"] 